/*
 * AqueousActivity.cpp
 *
 *  Created on: 23 May 2011
 *      Author: allan
 */

#include "AqueousActivity.h"

// C++ includes
#include <cmath>
#include <map>
#include <numeric>
using namespace std;

// GeoReact includes
#include "AqueousPhase.h"
#include "NumericalUtilities.h"

const double DebyeHuckelParameterA(double T, double P);
const double DebyeHuckelParameterB(double T, double P);
const double HKFParameterNaCl(double T, double P);
const double HKFParameterNapClm(double T, double P);

const DoubleVector EffectiveElectrostaticRadiiOfIons(const StringVector& ionsNames);

const double ElectrostaticParameterEta = 1.66027E+05;

//-------------------------------------------------------------------------------------------------//
const DoubleVector HkfIonicActivityCoefficients(const AqueousPhase& aqueousPhase, double T, double P, const DoubleVector& n, const DoubleVector& effRadii)
{
	// The electrical charges of the ionic species
	const DoubleVector zi = aqueousPhase.GetIonicCharges();
	
	// The number of ionic species
	const uint numIons = zi.size();
	
	// The effective ionic strength of the aqueousPhaseution
	const double I = aqueousPhase.EffectiveIonicStrength(n);
	
	// The square root of the ionic strength
	const double sqrtI = sqrt(I);
	
	// The total molality of the aqueousPhaseutes
	const double mT = aqueousPhase.TotalMolalitySolutes(n);
	
	// The parameters for the HKF model
	const double A       = DebyeHuckelParameterA(T, P);
	const double B       = DebyeHuckelParameterB(T, P);
	const double bNaCl   = HKFParameterNaCl(T, P);
	const double bNapClm = HKFParameterNapClm(T, P);
	
	// The activity coefficients of the ions
	DoubleVector gammai(numIons);
	
	// Loop over all ions in the aqueousPhaseution to compute their activity coefficients
	for(uint ion = 0; ion < numIons; ++ion)
	{
		// The electrical charge of the current ion
		const double z = zi[ion];

		// The square of the electrical charge of the ion
		const double z2 = z * z;

		// The effective electrostatic radius of the ion
		const double effRadius = effRadii[ion];

		// The abaqueousPhaseute Born coefficient of the ion
		const double omegaAbs = ElectrostaticParameterEta * z2 / effRadius;

		// The Debye-Huckel ion size parameter of the ion as computed by Reed, 1982 and also in TOUGHREACT
		const double a = (z < 0) ? 2.0 * (effRadius + 1.91 * abs(z))/(abs(z) + 1.0) : 
				                   2.0 * (effRadius + 1.81 * abs(z))/(abs(z) + 1.0);

		// The log10 of the activity coefficient of the current ion in molality scale
		const double log10_gi = -(A * z2 * sqrtI)/(1.0 + a * B * sqrtI) - log10(1.0 + 0.0180153 * mT) + (omegaAbs * bNaCl + bNapClm - 0.19 * (abs(z) - 1.0)) * I;
		
		// Store the activity coefficient of the current ion
		gammai[ion] = pow(10.0, log10_gi); 
	}
	
	return gammai;
}

const double HkfWaterActivityCoefficient(const AqueousPhase& aqueousPhase, double T, double P, const DoubleVector& n, const DoubleVector& effRadii)
{
	// The electrical charges of the ionic species
	const DoubleVector zi = aqueousPhase.GetIonicCharges();
	
	// The number of ionic species
	const uint numIons = zi.size();
	
	// The effective ionic strength of the aqueousPhaseution
	const double I = aqueousPhase.EffectiveIonicStrength(n);
	
	// The square root of the ionic strength
	const double sqrtI = sqrt(I);
	
	// The total molality of the aqueousPhaseutes
	const double mT = aqueousPhase.TotalMolalitySolutes(n);
	
	// The parameters for the HKF model
	const double A       = DebyeHuckelParameterA(T, P);
	const double B       = DebyeHuckelParameterB(T, P);
	const double bNaCl   = HKFParameterNaCl(T, P);
	const double bNapClm = HKFParameterNapClm(T, P);

	// The Born coefficient of the ion H[+]
	const double omegaH = 0.5387E+05;
	
	// The stoichiometric molality of the ions
	const DoubleVector smions = aqueousPhase.IonicStoichiometricMolalities(n);
	
	// The osmotic coefficient of the aqueous phase
	double osmoticCoeff = 0.0;
	
	// Loop over all ions in the aqueousPhaseution to compute the osmotic coefficient
	for(uint ion = 0; ion < numIons; ++ion)
	{
		// The electrical charge of the current ion species
		const double z = zi[ion];
		
		// The square of the electrical charge
		const double z2 = z * z;
		
		// The effective electrostatic radius of the current ion
		const double effRadius = effRadii[ion];
		
		// The Born coefficient of the current ion
		const double omega = ElectrostaticParameterEta * z2 / effRadius - z * omegaH;
		
		// The Debye-Huckel ion size parameter of the ion as computed by Reed, 1982 and also in TOUGHREACT
		const double a = (z < 0) ? 2.0 * (effRadius + 1.91 * abs(z))/(abs(z) + 1.0) : 
				                   2.0 * (effRadius + 1.81 * abs(z))/(abs(z) + 1.0);
		
		// The Lambda parameter of the current ion
		const double Lambda = 1.0 + a * B * sqrtI;
		
		// The sigma parameter of the current ion
		const double sigma = 3.0/pow(a * B * sqrtI, 3) * (Lambda - 1.0/Lambda - 2.0 * log(Lambda));
		
		// calculate the psi contribution of the current ion
		double psi = A * z2 * sqrtI * sigma/3.0 - log10(1.0 + 0.0180153 * mT)/(0.0180153 * mT) - 0.5 * (omega * bNaCl + bNapClm - 0.19 * (abs(z) - 1.0)) * I;
		
		osmoticCoeff += smions[ion] * psi;
	}
	
	// Correct the osmotic coefficient
	osmoticCoeff *= -2.303/mT;
	
	// Compute the activity of water
	const double activityWater = exp(-osmoticCoeff * mT/55.508);
	
	// Return the activity coefficient of water divided by 55.508, since we define gw = aw/55.508, where gw and aw are the activity coefficient and activity of water respectively
	return activityWater / 55.508;
}

const double SetschenowActivityCoefficient(const AqueousPhase& aqueousPhase, double T, double P, const DoubleVector& n, const DoubleVector& bi, const DoubleVector& Ti)
{
	// The mole-fraction of the water species to convert the activity coefficient from mole-fraction scale to molality scale
	const double xw = n[aqueousPhase.GetWaterIndex()] / std::accumulate(n.begin(), n.end(), 0.0);
	
	// The effective ionic strength of the aqueousPhaseution
	const double I = aqueousPhase.EffectiveIonicStrength(n);
	
	// The Setschenow constant at the given temperature
	const double b = LagrangeInterpolation(T, &Ti[0], &bi[0], Ti.size());
	
	// Return the activity coefficient of the neutral species in a molality scale
	return xw * pow(10.0, b * I);
}

const double DrummondCO2ActivityCoefficient(const AqueousPhase& aqueousPhase, double T, double P, const DoubleVector& n)
{
	// The temperature in degrees Kelvin
	const double Tk = T + 273.15;
	
	// The effective ionic strength of the aqueousPhaseution
	const double I = aqueousPhase.EffectiveIonicStrength(n);
	
	// Return the Drummond activity coefficient of aqueous CO2
	return exp((-1.0312 + 1.2806E-3*Tk + 255.9/Tk) * I - (0.4445 - 1.606E-3*Tk) * I/(I + 1.0));
}

const IonicActivityFunction HkfIonicActivityModel(const StringVector& ionsNames)
{
	const DoubleVector effRadii = EffectiveElectrostaticRadiiOfIons(ionsNames);
	
	return std::bind(HkfIonicActivityCoefficients, _1, _2, _3, _4, effRadii);
}

const NeutralActivityFunction SetschenowActivityModel()
{
	return SetschenowActivityModel(0.0);
}

const NeutralActivityFunction SetschenowActivityModel(double b)
{
	return SetschenowActivityModel({b}, {25.0}); // Assume the value b is given at 25.0 degC
}

const NeutralActivityFunction SetschenowActivityModel(const DoubleVector& bi, const DoubleVector& Ti)
{
	return std::bind(SetschenowActivityCoefficient, _1, _2, _3, _4, bi, Ti);
}

const NeutralActivityFunction HkfWaterActivityModel(const StringVector& ionsNames)
{
	const DoubleVector effRadii = EffectiveElectrostaticRadiiOfIons(ionsNames);
	
	return std::bind(HkfWaterActivityCoefficient, _1, _2, _3, _4, effRadii);
}

const NeutralActivityFunction DrummondCO2ActivityModel()
{
	return NeutralActivityFunction(DrummondCO2ActivityCoefficient);
}

//-------------------------------------------------------------------------------------------------//
std::map<string,double> effectiveElectrostaticRadius = 
{
	{"H[+]",    3.08},
	{"Li[+]",   1.64},
	{"Na[+]",   1.91},
	{"K[+]",    2.27},
	{"Rb[+]",   2.41},
	{"Cs[+]",   2.61},
	{"NH4[+]",  2.31},
	{"Ag[+]",   2.20},
	{"Au[+]",   2.31},
	{"Cu[+]",   1.90},
	{"Mg[2+]",  2.54},
	{"Sr[2+]",  3.00},
	{"Ca[2+]",  2.87},
	{"Ba[2+]",  3.22},
	{"Pb[2+]",  3.08},
	{"Zn[2+]",  2.62},
	{"Cu[2+]",  2.60},
	{"Cd[2+]",  2.85},
	{"Hg[2+]",  2.98},
	{"Fe[2+]",  2.62},
	{"Mn[2+]",  2.68},
	{"Fe[3+]",  3.46},
	{"Al[3+]",  3.33},
	{"Au[3+]",  3.72},
	{"La[3+]",  3.96},
	{"Gd[3+]",  3.79},
	{"In[3+]",  3.63},
	{"Ca[3+]",  3.44},
	{"F[-]",    1.33},
	{"Cl[-]",   1.81},
	{"Br[-]",   1.96},
	{"I[-]",    2.20},
	{"OH[-]",   1.40},
	{"HS[-]",   1.84},
	{"NO3[-]",  2.81},
	{"HCO3[-]", 2.10},
	{"HSO4[-]", 2.37},
	{"ClO4[-]", 3.59},
	{"ReO4[-]", 4.23},
	{"SO4[2-]", 3.15},
	{"CO3[2-]", 2.81}
};

const unsigned NumPointsOfDebyeHuckelA               = 8;
const unsigned NumPointsOfDebyeHuckelB               = 8;
const unsigned NumPointsOfDebyeHuckelParameterNaCl   = 12;
const unsigned NumPointsOfDebyeHuckelParameterNapClm = 12;

const double TemperatureOfDebyeHuckelA[]               = {0, 25.0, 60.0, 100.0, 150.0, 200.0, 250.0, 300.0};
const double TemperatureOfDebyeHuckelB[]               = {0, 25.0, 60.0, 100.0, 150.0, 200.0, 250.0, 300.0};
const double TemperatureOfDebyeHuckelParameterNaCl[]   = {25.0, 50.0, 75.0, 100.0, 125.0, 150.0, 175.0, 200.0, 225.0, 250.0, 275.0, 300.0};
const double TemperatureOfDebyeHuckelParameterNapClm[] = {25.0, 50.0, 75.0, 100.0, 125.0, 150.0, 175.0, 200.0, 225.0, 250.0, 275.0, 300.0};

const double ValuesOfDebyeHuckelA[]               = {0.4939, 0.5114, 0.5465, 0.5995, 0.6855, 0.7994, 0.9593, 1.2180};
const double ValuesOfDebyeHuckelB[]               = {0.3253, 0.3288, 0.3346, 0.6842, 0.3421, 0.3639, 0.3766, 0.3925};
const double ValuesOfDebyeHuckelParameterNaCl[]   = {2.47, 2.15, 1.79, 1.39, 0.93, 0.41, -0.18, -0.85, -1.64, -2.57, -3.71, -5.21};
const double ValuesOfDebyeHuckelParameterNapClm[] = {-9.77, -5.59, -2.43, 0.23, 2.44, 4.51, 6.38, 8.11, 9.73, 11.29, 12.71, 14.15}; 

const double DebyeHuckelParameterA(double T, double P)
{
	return LagrangeInterpolation(T, TemperatureOfDebyeHuckelA, ValuesOfDebyeHuckelA, NumPointsOfDebyeHuckelA);
}

const double DebyeHuckelParameterB(double T, double P)
{
	return LagrangeInterpolation(T, TemperatureOfDebyeHuckelB, ValuesOfDebyeHuckelB, NumPointsOfDebyeHuckelB);
}

const double HKFParameterNaCl(double T, double P)
{
	double R = 1.9858775;

	return 1.0E-3 * 1.0/(2.302585 * R * (T + 273.15)) * LagrangeInterpolation(T, TemperatureOfDebyeHuckelParameterNaCl, ValuesOfDebyeHuckelParameterNaCl, NumPointsOfDebyeHuckelParameterNaCl);
}

const double HKFParameterNapClm(double T, double P)
{
	return 1.0E-2 * LagrangeInterpolation(T, TemperatureOfDebyeHuckelParameterNapClm, ValuesOfDebyeHuckelParameterNapClm, NumPointsOfDebyeHuckelParameterNapClm);
}

const DoubleVector EffectiveElectrostaticRadiiOfIons(const StringVector& ionsNames)
{
	DoubleVector effRadii(ionsNames.size());
	
	// Loop over all the ions
	for(uint ion = 0; ion < ionsNames.size(); ++ion) 
	{
		auto iterator = effectiveElectrostaticRadius.find(ionsNames[ion]);
		
		// Check if the ion has been found in the list of ions
		if(iterator == effectiveElectrostaticRadius.end())
		{
			const int charge = (int)ElectricalCharge(ionsNames[ion]);
			
			// The approach below was taken from:
			// Xu, T., Sonnenthal, E., Spycher, N., & Pruess, K. (2004). TOUGHREACT-User's Guide: A Simulation Program for Non-isothermal Multiphase 
			// Reactive geochemical Transport in Variable Saturated Geologic Media. Computers & Geosciences (Vol. 32).
			if(charge == +4)
				effRadii[ion] = 4.5;
			else if(charge > +3)
				effRadii[ion] = std::abs(charge) * 4.5/4.0;
			else // charge < -3
				effRadii[ion] = std::abs(charge) * 4.2/3.0;
		}
		
		effRadii[ion] = iterator->second;
	}
	
	return effRadii;
}
